# 导入操作系统相关模块，用于文件和目录操作
import os
# 导入 PyTorch 深度学习框架
import torch
# 导入 NumPy 科学计算库
import numpy as np
# 从 typing 模块导入 List 类型注解
from typing import List
# 从 torch 模块导入 Tensor 类型
from torch import Tensor
# 从 multiprocessing 模块导入 Pipe 和 Process，用于多进程通信和进程创建
from multiprocessing import Pipe, Process


# 定义 Config 类，用于存储和管理强化学习训练过程中的各种配置参数
class Config:
    # 类的构造函数，初始化各种配置参数
    def __init__(self, agent_class=None, env_class=None, env_args=None):
        # 子环境的数量，初始化为 None
        self.num_envs = None
        # 智能体类，用于创建智能体实例
        self.agent_class = agent_class  # agent = agent_class(...)
        # 判断是否为离线策略算法
        self.if_off_policy = self.get_if_off_policy()  # whether off-policy or on-policy of DRL algorithm

        '''环境相关参数'''
        # 环境类，用于创建环境实例
        self.env_class = env_class  # env = env_class(**env_args)
        # 环境参数，用于配置环境实例
        self.env_args = env_args  # env = env_class(**env_args)
        # 如果没有提供环境参数，则使用默认的环境参数
        if env_args is None:  # dummy env_args
            env_args = {'env_name': None,
                        'num_envs': 1,
                        'max_step': 12345,
                        'state_dim': None,
                        'action_dim': None,
                        'if_discrete': None, }
        # 设置默认的子环境数量为 1
        env_args.setdefault('num_envs', 1)  # `num_envs=1` in default in single env.
        # 设置默认的最大步数为 12345，这是一个足够大的值
        env_args.setdefault('max_step', 12345)  # `max_step=12345` in default, which is a large enough value.
        # 环境名称，用于设置模型保存的工作目录
        self.env_name = env_args['env_name']  # the name of environment. Be used to set 'cwd'.
        # 子环境的数量，单环境时默认为 1
        self.num_envs = env_args['num_envs']  # the number of sub envs in vectorized env. `num_envs=1` in single env.
        # 一个回合的最大步数，默认设置为 12345
        self.max_step = env_args['max_step']  # the max step number of an episode. 'set as 12345 in default.
        # 状态的向量维度（特征数量）
        self.state_dim = env_args['state_dim']  # vector dimension (feature number) of state
        # 动作的向量维度（特征数量）
        self.action_dim = env_args['action_dim']  # vector dimension (feature number) of action
        # 动作空间是离散还是连续
        self.if_discrete = env_args['if_discrete']  # discrete or continuous action space

        '''奖励塑形相关参数'''
        # 未来奖励的折扣因子
        self.gamma = 0.99  # discount factor of future rewards
        # 奖励缩放因子，近似目标奖励通常接近 256
        self.reward_scale = 2 ** 0  # an approximate target reward usually be closed to 256

        '''训练相关参数'''
        # 多层感知机（MLP）的中间层维度
        self.net_dims = (64, 32)  # the middle layer dimension of MLP (MultiLayer Perceptron)
        # 网络更新的学习率
        self.learning_rate = 6e-5  # the learning rate for network updating
        # 梯度裁剪的范数，范围在 0.1 ~ 4.0 之间，用于在归一化后裁剪梯度
        self.clip_grad_norm = 3.0  # 0.1 ~ 4.0, clip the gradient after normalization
        # 值和状态归一化的 tau 参数，`std = (1-std)*std + tau*std`
        self.state_value_tau = 0  # the tau of normalize for value and state `std = (1-std)*std + tau*std`
        # 软目标更新的 tau 参数，2 ** -8 约等于 5e-3，`net = (1-tau)*net + tau*net1`
        self.soft_update_tau = 5e-3  # 2 ** -8 ~= 5e-3. the tau of soft target update `net = (1-tau)*net + tau*net1`
        # 如果是离线策略算法
        if self.if_off_policy:  # off-policy
            # 从经验回放缓冲区中采样的转移样本数量
            self.batch_size = int(64)  # num of transitions sampled from replay buffer.
            # 探索时收集的步数，收集完成后更新网络
            self.horizon_len = int(512)  # collect horizon_len step while exploring, then update networks
            # 经验回放缓冲区的大小，离线策略算法使用先进先出策略
            self.buffer_size = int(1e6)  # ReplayBuffer size. First in first out for off-policy.
            # 重复更新网络的次数，用于保持评论家网络损失较小
            self.repeat_times = 1.0  # repeatedly update network using ReplayBuffer to keep critic's loss small
            # 是否使用优先经验回放（PER）来处理稀疏奖励
            self.if_use_per = False  # use PER (Prioritized Experience Replay) for sparse reward
        else:  # 在线策略算法
            # 从经验回放缓冲区中采样的转移样本数量
            self.batch_size = int(128)  # num of transitions sampled from replay buffer.
            # 探索时收集的步数，收集完成后更新网络
            self.horizon_len = int(2048)  # collect horizon_len step while exploring, then update network
            # 经验回放缓冲区的大小，在线策略算法清空经验回放缓冲区
            self.buffer_size = None  # ReplayBuffer size. Empty the ReplayBuffer for on-policy.
            # 重复更新网络的次数，用于保持评论家网络损失较小
            self.repeat_times = 8.0  # repeatedly update network using ReplayBuffer to keep critic's loss small
            # 是否使用 V-trace + GAE（广义优势估计）来处理稀疏奖励
            self.if_use_vtrace = False  # use V-trace + GAE (Generalized Advantage Estimation) for sparse reward

        '''设备相关参数'''
        # 单个 GPU 的 ID，-1 表示使用 CPU
        self.gpu_id = int(0)  # `int` means the ID of single GPU, -1 means CPU
        # 每个 GPU 的 rollout 工作进程数量（可调整以提高 GPU 利用率）
        self.num_workers = 2  # rollout workers number pre GPU (adjust it to get high GPU usage)
        # PyTorch 使用的 CPU 线程数，`torch.set_num_threads(self.num_threads)`
        self.num_threads = 8  # cpu_num for pytorch, `torch.set_num_threads(self.num_threads)`
        # 随机种子，在训练前初始化
        self.random_seed = 0  # initialize random seed in self.init_before_training()
        # 学习器使用的 GPU ID，-1 表示使用 CPU
        self.learner_gpus = 0  # `int` means the ID of single GPU, -1 means CPU

        '''评估相关参数'''
        # 当前工作目录，用于保存模型。None 表示自动设置
        self.cwd = None  # current working directory to save model. None means set automatically
        # 是否移除当前工作目录文件夹？（True: 移除，False: 保留，None: 询问用户）
        self.if_remove = True  # remove the cwd folder? (True, False, None:ask me)
        # 如果总步数超过该值，则停止训练
        self.break_step = np.inf  # break training if 'total_step > break_step'
        # 如果累计奖励超过该值，则停止训练
        self.break_score = np.inf  # break training if `cumulative_rewards > break_score`
        # 是否持续保存检查点。False 表示只在训练停止时保存
        self.if_keep_save = True  # keeping save the checkpoint. False means save until stop training.
        # 是否覆盖最佳策略网络，保存路径为 `self.cwd/actor.pth`
        self.if_over_write = False  # overwrite the best policy network. `self.cwd/actor.pth`
        # 是否保存经验回放缓冲区，以便在训练停止后继续训练
        self.if_save_buffer = False  # if save the replay buffer for continuous training after stop training

        # 保存演员网络的间隔步数，保存路径为 f"{cwd}/actor_*.pth"，用于绘制学习曲线
        self.save_gap = int(8)  # save actor f"{cwd}/actor_*.pth" for learning curve.
        # 获取平均回合累计回报的评估次数
        self.eval_times = int(3)  # number of times that get the average episodic cumulative return
        # 每训练多少步评估一次智能体
        self.eval_per_step = int(2e4)  # evaluate the agent per training steps
        # 评估环境类，用于创建评估环境实例
        self.eval_env_class = None  # eval_env = eval_env_class(*eval_env_args)
        # 评估环境参数，用于配置评估环境实例
        self.eval_env_args = None  # eval_env = eval_env_class(*eval_env_args)

    # 训练前的初始化操作
    def init_before_training(self):
        # 设置 NumPy 的随机种子
        np.random.seed(self.random_seed)
        # 设置 PyTorch 的随机种子
        torch.manual_seed(self.random_seed)
        # 设置 PyTorch 使用的 CPU 线程数
        torch.set_num_threads(self.num_threads)
        # 设置 PyTorch 的默认数据类型为 float32
        torch.set_default_dtype(torch.float32)

        '''设置用于保存模型的当前工作目录'''
        # 如果未指定当前工作目录，则自动设置
        if self.cwd is None:  # set cwd (current working directory) for saving model
            self.cwd = f'./{self.env_name}_{self.agent_class.__name__[5:]}_{self.random_seed}'

        '''移除历史文件'''
        # 如果 if_remove 为 None，则询问用户是否移除当前工作目录
        if self.if_remove is None:
            self.if_remove = bool(input(f"| Arguments PRESS 'y' to REMOVE: {self.cwd}? ") == 'y')
        # 如果 if_remove 为 True，则移除当前工作目录
        if self.if_remove:
            import shutil
            shutil.rmtree(self.cwd, ignore_errors=True)
            print(f"| Arguments Remove cwd: {self.cwd}")
        else:
            print(f"| Arguments Keep cwd: {self.cwd}")
        # 创建当前工作目录，如果已存在则忽略错误
        os.makedirs(self.cwd, exist_ok=True)

    # 判断是否为离线策略算法
    def get_if_off_policy(self) -> bool:
        # 获取智能体类的名称，如果未提供则为空字符串
        agent_name = self.agent_class.__name__ if self.agent_class else ''
        # 在线策略算法的名称列表
        on_policy_names = ('SARSA', 'VPG', 'A2C', 'A3C', 'TRPO', 'PPO', 'MPO')
        # 如果智能体名称不包含任何在线策略算法的名称，则返回 True，表示是离线策略算法
        return all([agent_name.find(s) == -1 for s in on_policy_names])

    # 打印配置参数
    def print(self):
        from pprint import pprint
        pprint(vars(self))  # prints out args in a neat, readable format


# 构建环境实例的函数
def build_env(env_class=None, env_args: dict = None, gpu_id: int = -1):
    # 在构建向量化环境前设置 GPU ID
    env_args['gpu_id'] = gpu_id  # set gpu_id for vectorized env before build it

    # 如果需要构建向量化环境
    if env_args.get('if_build_vec_env'):
        # 获取子环境的数量
        num_envs = env_args['num_envs']
        # 创建向量化环境实例
        env = VecEnv(env_class=env_class, env_args=env_args, num_envs=num_envs, gpu_id=gpu_id)
    # 如果是标准的 OpenAI Gym 环境
    elif env_class.__module__ == 'gym.envs.registration':
        import gym
        # 确保 Gym 版本在指定范围内
        assert '0.18.0' <= gym.__version__ <= '0.25.2'  # pip3 install gym==0.24.0
        # 设置 Gym 的日志级别，屏蔽警告信息
        gym.logger.set_level(40)  # Block warning
        # 创建 Gym 环境实例
        env = env_class(id=env_args['env_name'])
    else:
        # 使用过滤后的参数创建自定义环境实例
        env = env_class(**kwargs_filter(env_class.__init__, env_args.copy()))

    # 设置默认的子环境数量为 1
    env_args.setdefault('num_envs', 1)
    # 设置默认的最大步数为 12345
    env_args.setdefault('max_step', 12345)

    # 将环境参数中的属性设置到环境实例上
    for attr_str in ('env_name', 'num_envs', 'max_step', 'state_dim', 'action_dim', 'if_discrete'):
        setattr(env, attr_str, env_args[attr_str])
    return env


# 过滤函数参数的函数，只保留函数签名中包含的参数
def kwargs_filter(function, kwargs: dict) -> dict:
    import inspect
    # 获取函数签名的参数
    sign = inspect.signature(function).parameters.values()
    # 将参数名提取到集合中
    sign = {val.name for val in sign}
    # 获取函数签名参数和传入参数的交集
    common_args = sign.intersection(kwargs.keys())
    # 返回过滤后的参数
    return {key: kwargs[key] for key in common_args}  # filtered kwargs


# 获取标准 OpenAI Gym 环境信息的函数
def get_gym_env_args(env, if_print: bool) -> dict:
    """get a dict about a standard OpenAI gym env information.
    assert 0.18.0 <= gym.__version__ <= 0.25.3

    env: a standard OpenAI gym env
    if_print: [bool] print the dict about env information.
    return: env_args [dict]

    env_args = {
        'env_name': env_name,       # [str] the environment name, such as XxxXxx-v0
        'num_envs': num_envs.       # [int] the number of sub envs in vectorized env. `num_envs=1` in single env.
        'max_step': max_step,       # [int] the max step number of an episode.
        'state_dim': state_dim,     # [int] the dimension of state
        'action_dim': action_dim,   # [int] the dimension of action or the number of discrete action
        'if_discrete': if_discrete, # [bool] action space is discrete or continuous
    }
    """
    import gym

    # 判断是否为标准的 OpenAI Gym 环境
    if_gym_standard_env = {'unwrapped', 'observation_space', 'action_space', 'spec'}.issubset(dir(env))

    # 如果是标准的 OpenAI Gym 环境且没有 num_envs 属性
    if if_gym_standard_env and (not hasattr(env, 'num_envs')):  # isinstance(env, gym.Env):
        # 确保 Gym 版本在指定范围内
        assert '0.18.0' <= gym.__version__ <= '0.25.2'  # pip3 install gym==0.24.0
        # 获取环境的名称
        env_name = env.unwrapped.spec.id
        # 获取子环境的数量，默认为 1
        num_envs = getattr(env, 'num_envs', 1)
        # 获取一个回合的最大步数，默认为 12345
        max_step = getattr(env, '_max_episode_steps', 12345)

        # 获取状态空间的形状
        state_shape = env.observation_space.shape
        # 如果状态空间是一维的，则取其第一个维度作为状态维度，否则直接使用状态形状
        state_dim = state_shape[0] if len(state_shape) == 1 else state_shape  # sometimes state_dim is a list

        # 判断动作空间是否为离散空间
        if_discrete = isinstance(env.action_space, gym.spaces.Discrete)
        # 如果是离散动作空间
        if if_discrete:  # make sure it is discrete action space
            # 获取离散动作的数量
            action_dim = getattr(env.action_space, 'n')
        # 如果是连续动作空间
        elif isinstance(env.action_space, gym.spaces.Box):  # make sure it is continuous action space
            # 获取连续动作的维度
            action_dim = env.action_space.shape[0]
            # 如果动作空间的上限不等于 1，则打印警告信息
            if any(env.action_space.high - 1):
                print('WARNING: env.action_space.high', env.action_space.high)
            # 如果动作空间的下限不等于 -1，则打印警告信息
            if any(env.action_space.low + 1):
                print('WARNING: env.action_space.low', env.action_space.low)
        else:
            raise RuntimeError('\n| Error in get_gym_env_info(). Please set these value manually:'
                               '\n  `state_dim=int; action_dim=int; if_discrete=bool;`'
                               '\n  And keep action_space in range (-1, 1).')
    else:
        # 获取环境名称，默认为 'env'
        env_name = getattr(env, 'env_name', 'env')
        # 获取子环境的数量，默认为 1
        num_envs = getattr(env, 'num_envs', 1)
        # 获取一个回合的最大步数，默认为 12345
        max_step = getattr(env, 'max_step', 12345)
        # 获取状态维度
        state_dim = env.state_dim
        # 获取动作维度
        action_dim = env.action_dim
        # 获取动作空间是否为离散空间
        if_discrete = env.if_discrete

    # 构建环境参数字典
    env_args = {'env_name': env_name,
                'num_envs': num_envs,
                'max_step': max_step,
                'state_dim': state_dim,
                'action_dim': action_dim,
                'if_discrete': if_discrete, }
    # 如果需要打印环境参数
    if if_print:
        # 格式化环境参数字符串，使其更易读
        env_args_str = repr(env_args).replace(',', f",\n{'':11}")
        print(f"env_args = {env_args_str}")
    return env_args


"""向量化环境"""


# 定义 SubEnv 类，继承自 Process 类，用于创建子环境进程
class SubEnv(Process):
    # 类的构造函数，初始化子环境进程的参数
    def __init__(self, sub_pipe0: Pipe, vec_pipe1: Pipe,
                 env_class, env_args: dict, env_id: int = 0):
        super().__init__()
        # 子进程接收动作的管道
        self.sub_pipe0 = sub_pipe0
        # 子进程发送结果的管道
        self.vec_pipe1 = vec_pipe1

        # 环境类，用于创建环境实例
        self.env_class = env_class
        # 环境参数，用于配置环境实例
        self.env_args = env_args
        # 子环境的 ID
        self.env_id = env_id

    # 子进程的运行函数
    def run(self):
        # 关闭 PyTorch 的梯度计算，提高性能
        torch.set_grad_enabled(False)

        '''构建环境'''
        # 如果是标准的 OpenAI Gym 环境
        if self.env_class.__module__ == 'gym.envs.registration':  # is standard OpenAI Gym env
            # 创建 Gym 环境实例
            env = self.env_class(id=self.env_args['env_name'])
        else:
            # 使用过滤后的参数创建自定义环境实例
            env = self.env_class(**kwargs_filter(self.env_class.__init__, self.env_args.copy()))

        '''设置环境的随机种子'''
        # 使用子环境 ID 作为随机种子
        random_seed = self.env_id
        # 设置 NumPy 的随机种子
        np.random.seed(random_seed)
        # 设置 PyTorch 的随机种子
        torch.manual_seed(random_seed)

        # 进入无限循环，不断接收动作并执行环境交互
        while True:
            # 从管道接收动作
            action = self.sub_pipe0.recv()
            # 如果接收到的动作是 None，则重置环境
            if action is None:
                state = env.reset()
                # 将子环境 ID 和重置后的状态发送到管道
                self.vec_pipe1.send((self.env_id, state))
            else:
                # 执行环境的 step 操作，获取状态、奖励、是否结束和额外信息
                state, reward, done, info_dict = env.step(action)
                # 如果回合结束，则重置环境
                state = env.reset() if done else state
                # 将子环境 ID、状态、奖励、是否结束和额外信息发送到管道
                self.vec_pipe1.send((self.env_id, state, reward, done, info_dict))


# 定义 VecEnv 类，用于创建向量化环境
class VecEnv:
    # 类的构造函数，初始化向量化环境的参数
    def __init__(self, env_class: object, env_args: dict, num_envs: int, gpu_id: int = -1):
        # 根据 GPU ID 创建 PyTorch 设备对象，如果 GPU 可用且 ID 大于等于 0，则使用 GPU，否则使用 CPU
        self.device = torch.device(f"cuda:{gpu_id}" if (torch.cuda.is_available() and (gpu_id >= 0)) else "cpu")
        # 向量化环境中包含的子环境数量
        self.num_envs = num_envs  # the number of sub env in vectorized env.

        '''设计自定义环境时所需的必要环境信息'''
        # 环境的名称
        self.env_name = env_args['env_name']  # the name of this env.
        # 一个回合的最大步数，用于评估
        self.max_step = env_args['max_step']  # the max step number in an episode for evaluation
        # 状态的特征数量
        self.state_dim = env_args['state_dim']  # feature number of state
        # 动作的特征数量
        self.action_dim = env_args['action_dim']  # feature number of action
        # 动作空间是离散还是连续
        self.if_discrete = env_args['if_discrete']  # discrete action or continuous action

        '''使用多进程加速：Process, Pipe'''
        # 确保子环境数量不超过 64
        assert self.num_envs <= 64
        # 用于存储每个子环境返回结果的列表
        self.res_list = [[] for _ in range(self.num_envs)]

        # 创建 num_envs 对单向管道，用于子进程和主进程之间的通信
        sub_pipe0s, sub_pipe1s = list(zip(*[Pipe(duplex=False) for _ in range(self.num_envs)]))
        # 保存子进程发送动作的管道
        self.sub_pipe1s = sub_pipe1s

        # 创建一对单向管道，用于接收子进程的返回结果
        vec_pipe0, vec_pipe1 = Pipe(duplex=False)  # recv, send
        # 保存接收结果的管道
        self.vec_pipe0 = vec_pipe0

        # 创建 num_envs 个子环境进程
        self.sub_envs = [
            SubEnv(sub_pipe0=sub_pipe0, vec_pipe1=vec_pipe1,
                   env_class=env_class, env_args=env_args, env_id=env_id)
            for env_id, sub_pipe0 in enumerate(sub_pipe0s)
        ]

        # 将所有子进程设置为守护进程，确保主进程退出时子进程也能安全退出
        [setattr(p, 'daemon', True) for p in self.sub_envs]  # set before process start to exit safely
        # 启动所有子进程
        [p.start() for p in self.sub_envs]

    # 重置向量化环境中的所有子环境
    def reset(self) -> Tensor:  # reset the agent in env
        # 关闭 PyTorch 的梯度计算，提高性能
        torch.set_grad_enabled(False)

        # 向每个子进程发送 None 信号，通知其重置环境
        for pipe in self.sub_pipe1s:
            pipe.send(None)
        # 获取有序的子环境重置状态
        states, = self.get_orderly_zip_list_return()
        # 将状态转换为 PyTorch 张量，并移动到指定设备上
        states = torch.tensor(np.stack(states), dtype=torch.float32, device=self.device)
        return states

    # 向量化环境执行一步交互操作
    def step(self, action: Tensor) -> (Tensor, Tensor, Tensor, List[dict]):  # agent interacts in env
        # 将动作张量从 GPU 移动到 CPU 并转换为 NumPy 数组
        action = action.detach().cpu().numpy()
        # 如果是离散动作空间，去掉多余的维度
        if self.if_discrete:
            action = action.squeeze(1)
        # 将动作发送到每个子进程
        for pipe, a in zip(self.sub_pipe1s, action):
            pipe.send(a)

        # 获取有序的子环境交互结果
        states, rewards, dones, info_dicts = self.get_orderly_zip_list_return()
        # 将状态转换为 PyTorch 张量，并移动到指定设备上
        states = torch.tensor(np.stack(states), dtype=torch.float32, device=self.device)
        # 将奖励转换为 PyTorch 张量，并移动到指定设备上
        rewards = torch.tensor(rewards, dtype=torch.float32, device=self.device)
        # 将是否结束标志转换为 PyTorch 张量，并移动到指定设备上
        dones = torch.tensor(dones, dtype=torch.bool, device=self.device)
        return states, rewards, dones, info_dicts

    # 关闭向量化环境，终止所有子进程
    def close(self):
        [process.terminate() for process in self.sub_envs]

    # 获取有序的子环境返回结果
    def get_orderly_zip_list_return(self):
        # 循环 num_envs 次，接收每个子进程的返回结果
        for _ in range(self.num_envs):
            res = self.vec_pipe0.recv()
            # 将结果存储到对应的子环境位置
            self.res_list[res[0]] = res[1:]
        # 将结果按子环境顺序重新排列并返回
        return list(zip(*self.res_list))
